from __future__ import division
from __future__ import print_function

import copy
import numpy as np
import pandas as pd

from qlib.data import D

def risk_analysis(r, method='ci'):
    """Risk Analysis

    Parameters
    ----------
    r : pandas.Series
        daily return series
    method : str
        interest calculation method, ci(compound interest)/si(simple interest)
    """
    mean = r.mean()
    std = r.std(ddof=1)
    annual = (1+mean)**250-1 if method == 'ci' else mean*250
    sharpe = annual/std/np.sqrt(250)
    mdd = ((r.cumsum()-r.cumsum().cummax())/(1+r.cumsum().cummax())).min()
    data = {'mean': mean, 'std': std, 'annual': annual, 'sharpe': sharpe, 'mdd': mdd}
    res = pd.Series(data, index=data.keys()).to_frame('risk')
    return res

def backtest_weighted_portfolio(pred, account=1e9, risk_degree=0.80,
             open_cost=0.0015, close_cost=0.0025, benchmark='SH000905', shift=1, verbose=True):
    """BackTest

    Parameters
    ----------
    pred : pandas.DataFrame
        predict should has <instrument, datetime> index and one `score` column
    account : float
        init account value
    risk_degree : float
        position percentage of total value
    open_cost : float
        open transaction cost
    close_cost : float
        close transaction cost
    benchmark : str
        benchmark code
    shift : int
        whether to shift prediction by one day
    verbose : bool
        whether to print log
    """

    codes = sorted(pred.index.get_level_values(0).unique())
    dates = sorted(pred.index.get_level_values(1).unique())
    quote = D.features(codes, ['$vwap','$close','$change','$open'], dates[0], dates[-1], disk_cache=True).dropna()
    if benchmark is not None:
        bench = D.features([benchmark], ['$change'], dates[0], dates[-1], disk_cache=True).loc[benchmark, '$change']
    else:
        bench = quote.reindex(pred.index).groupby(level=1)['$change'].mean()
    quote = quote.to_dict('index')
    trade_dates = sorted(bench.index)
    
    returns = dict()
    costs = dict()
    turnovers = dict()
    accounts = dict()
    positions = dict()
    values = dict()

    w = dict()
    cash = account
    for pdate, date in zip(trade_dates, trade_dates[shift:]):
        if verbose:
            print('[I {:%Y-%m-%d}]: trade begin.'.format(date))
        # load target position
        ref = pred.loc(axis=0)[:, pdate] # today
        new = ref[ref['score'] != 0].reset_index(level=1, drop=True)['score'].to_dict()
        # initialize
        rtn = ct = to = val = 0 # return/cost/turnover/value
        # handle positions
        plan = dict()
        for code in set(list(new)+list(w)):
            # suspend
            if (code, date) not in quote or (code, pdate) not in quote:
                
                if code in w:
                    if verbose:
                        print('[W {:%Y-%m-%d}]: {} suspend.'.format(date, code))
                    val += w[code]['price'] * w[code]['amount']
                elif code not in w:
                    if verbose:
                        print('[W {:%Y-%m-%d}]: {} suspend.'.format(date, code))
                continue
            # trade
            amount = account * risk_degree * new.get(code, 0) / quote[(code, pdate)]['$close'] // 100 * 100
            price = quote[(code, date)]['$vwap']
            action = 'keep'
            if code not in w:
                # buy
                if amount == 0:
                    continue
                action = 'buy'
                trade_val = price * amount
                cost = trade_val * open_cost
                cash -= trade_val + cost
                profit = 0
                w[code] = {'price': price, 'amount': amount}
            elif w[code]['amount'] < amount:
                # buy
                action = 'buy'
                amount = amount - w[code]['amount']
                trade_val = price * amount
                cost = trade_val * open_cost
                cash -= trade_val + cost
                profit = 0
                w[code]['price'] = (w[code]['price']*w[code]['amount'] + trade_val) / (w[code]['amount'] + amount)
                w[code]['amount'] += amount
            elif w[code]['amount'] > amount:
                # sell
                action = 'sell'
                amount = w[code]['amount'] - amount
                trade_val = price * amount
                profit = trade_val - w[code]['price'] * amount
                cost = trade_val * close_cost
                cash += trade_val - cost
                if w[code]['amount'] == amount:
                    del w[code]
                    continue
                w[code]['price'] = (w[code]['price']*w[code]['amount'] - trade_val) / (w[code]['amount'] - amount)
                w[code]['amount'] -= amount
            # settle
            profit += (quote[(code, date)]['$close'] - w[code]['price']) * w[code]['amount']
            w[code]['price'] = quote[(code, date)]['$close']
            val += w[code]['price'] * w[code]['amount']
            ct += cost
            to += trade_val
            rtn += profit
            if verbose:
                print('[I {:%Y-%m-%d}]: {} {}, price {:.2f}, amount {}, value {:.2f}, earn {:.2f}.'.format(
                    date, action, code, price, amount, trade_val, profit - cost))
        # settle
        last_account = account
        account += rtn - ct
        rtn /= last_account
        to /= last_account
        ct /= last_account
        if verbose:
            print('[I {:%Y-%m-%d}]: trade end, account {:.2f}, profit {:.2f}, current cash {:.2f}.'.format(
                date, account, account - last_account, cash))
        # update
        returns[date] = rtn
        turnovers[date] = to
        costs[date] = ct
        accounts[date] = account
        positions[date] = copy.deepcopy(w)
        values[date] = val
        
    # create report
    report = pd.DataFrame()
    report['account'] = pd.Series(accounts)
    report['return'] = pd.Series(returns)
    report['turnover'] = pd.Series(turnovers)
    report['cost'] = pd.Series(costs)
    report['value'] = pd.Series(values)
    report['bench'] = bench
    report.index.name = 'date'
    
    return report, positions